import os
import yaml
import torch
import numpy as np
from tqdm import tqdm
from sklearn.metrics import classification_report, confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt
from data_preparation import prepare_dataloaders
from train import EmotionClassifier
from llava.model import LlavaLlamaForCausalLM

# 获取当前文件的绝对路径
current_dir = os.path.dirname(os.path.abspath(__file__))
# 获取项目根目录
root_dir = os.path.dirname(current_dir)

def evaluate(config, model, test_loader, device):
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for batch in tqdm(test_loader, desc="Testing"):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            pixel_values = batch['pixel_values'].to(device)
            labels = batch['labels'].to(device)
            
            outputs = model(input_ids, attention_mask, pixel_values)
            _, predicted = outputs.max(1)
            
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    # 计算分类报告
    emotion_labels = ['amusement', 'anger', 'awe', 'contentment', 'disgust', 'excitement', 'fear', 'sadness']
    report = classification_report(all_labels, all_preds, target_names=emotion_labels, digits=4)
    
    # 绘制混淆矩阵
    cm = confusion_matrix(all_labels, all_preds)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=emotion_labels, yticklabels=emotion_labels)
    plt.title('Confusion Matrix')
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.tight_layout()
    
    # 保存结果
    results_dir = config['output']['save_dir']
    os.makedirs(results_dir, exist_ok=True)
    
    # 保存分类报告
    with open(os.path.join(results_dir, 'classification_report.txt'), 'w') as f:
        f.write(report)
    
    # 保存混淆矩阵图
    plt.savefig(os.path.join(results_dir, 'confusion_matrix.png'))
    
    return report

if __name__ == '__main__':
    # 加载配置
    config_path = os.path.join(root_dir, 'configs', 'config.yaml')
    print(f"正在加载配置文件: {config_path}")
    
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    
    # 设置设备
    device = torch.device(config['model']['device'])
    
    # 加载数据
    _, _, test_loader = prepare_dataloaders(config)
    
    # 加载预训练模型
    llava_model = LlavaLlamaForCausalLM.from_pretrained(
        config['model']['pretrained_path'],
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True
    )
    
    # 创建分类器并加载训练好的权重
    model = EmotionClassifier(llava_model)
    model.load_state_dict(torch.load(os.path.join(config['output']['model_save_dir'], 'best_model.pth')))
    model = model.to(device)
    
    # 评估模型
    report = evaluate(config, model, test_loader, device)
    print("\n分类报告:")
    print(report) 